#! /usr/bin/python
import os
import argparse

COMMENT_HEADER = """
# if want more parameters, please visit the image url
# image url: https://hub.docker.com/_/zookeeper
# docker pull docker pull zookeeper
"""

VERSION = "3.1"

CONF_DIR = './conf'

STORAGE_DIR = "./storage"

DEFAULT_SERVICE_CONF = {
    'node_prefix': 'zk',
    'zoo_cfg': CONF_DIR + '/zoo.cfg',
    'storage_dir': STORAGE_DIR,
    'data_dir': STORAGE_DIR + '/data',
    'datalog_dir': STORAGE_DIR + '/datalog',
    'log_dir': STORAGE_DIR + '/logs'
}

SERVICE_TEMPLATE = """  {node_prefix}{num}:
    image: zookeeper
    restart: always
    hostname: {node_prefix}{num}
    ports:
      - {port}:2181
    environment:
      ZOO_MY_ID: {num}
    volumes:
      - "{zoo_cfg}:/conf/zoo.cfg"
      - "{data_dir}/{node_prefix}{num}:/data"
      - "{datalog_dir}/{node_prefix}{num}:/datalog"
      - "{log_dir}/{node_prefix}{num}:/logs"
"""

DOCKER_COMPOSE_TEMPLATE = """
{comment}

version: "{version}"
services:
{services}
"""

ZOO_CONF_TEMPLATE = """
dataDir=/data
dataLogDir=/datalog
tickTime=2000
initLimit=5
syncLimit=2
autopurge.snapRetainCount=3
autopurge.purgeInterval=0
maxClientCnxns=60
standaloneEnabled=true
admin.enableServer=true

{hosts}
"""

HOST_TEMPLATE = "server.{num}={node_prefix}{num}:2888:3888;2181"

def run(args):
    services = []
    hosts = []
    num = int(args.num)
    for i in range(num):
        conf = dict(DEFAULT_SERVICE_CONF)
        conf.update({
            'num' : i + 1,
            'node_prefix' : args.node_prefix,
            'port' : i + int(args.start_port)
        })
        services.append(SERVICE_TEMPLATE.format(**conf))
        hosts.append(HOST_TEMPLATE.format(**conf))

    services_str = ''.join(services)
    zoo_host_str = '\n'.join(hosts)

    docker_compose_str = DOCKER_COMPOSE_TEMPLATE.format(**{
        'comment' : COMMENT_HEADER,
        'version' : VERSION,
        'services' : services_str
    })

    zoo_cfg_str = ZOO_CONF_TEMPLATE.format(hosts =zoo_host_str)

    with open(args.output, 'w') as dcf:
        dcf.write(docker_compose_str)
        print("docker-compose.yml is saved in " + os.path.abspath(args.output))

    os.makedirs(CONF_DIR, exist_ok=True)
    zoo_cfg_path = os.path.abspath(os.path.join(CONF_DIR, "zoo.cfg"))
    with open(zoo_cfg_path, "w") as zcf:
        zcf.write(zoo_cfg_str)
        print("zoo.yml is saved in " + zoo_cfg_path)
        

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("-n", "--num", help="the number of service, default 3", default=3)
    parser.add_argument("-sp", "--start_port", help="the start port using by service, default 2181", default=2181)
    parser.add_argument("-np", "--node_prefix", help="the prefix of each service, default zk", default='zk')
    parser.add_argument("-o", "--output", help="output path of docker-compose.yml, default ./docker-compose.yml", default="docker-compose.yml")

    args = parser.parse_args()
    run(args)

if __name__ == '__main__':
    main()